/*
Xern Entity Relocation Network
Copyright (C) 2011 Ronie Salgado <roniesalg@gmail.com>

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
*/

#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <boost/make_shared.hpp>
#include "GpuRayTracer.hpp"

namespace XernRay
{
    struct Ray
    {
        Vector3 origin;
        Vector3 direction;
    };

    struct GpuFace
    {
        Plane plane;
    };

    struct GpuLight
    {
        Vector3 position;
        Vector3 direction;
        //Vector3 lightColor;
    };

    RayTracer::RayTracer()
    {
        passthroughProgram = Program::CreateFromFile("passthrough.glslv", "passthrough.glslf");
        primaryCastProgram = Program::CreateFromFile("primary.glslv", "primary.glslf");
        shadeProgram = Program::CreateFromFile("shade.glslv", "shade.glslf");
        width = 640;
        height = 480;
        valid = false;
    }

    RayTracer::~RayTracer()
    {
    }

    struct FaceCounter
    {
        FaceCounter(size_t &count)
            : count(count) {}

        void operator()(boost::shared_ptr<Object> object)
        {
            if(object->brush)
                count += object->brush->faces.size();
        }

        size_t& count;
    };

    struct BrushCounter
    {
        BrushCounter(size_t &count)
            : count(count) {}

        void operator()(boost::shared_ptr<Object> object)
        {
            if(object->brush)
                count++;
        }

        size_t& count;
    };

    struct LightCounter
    {
        LightCounter(size_t &count)
            : count(count) {}

        void operator()(boost::shared_ptr<Object> object)
        {
            if(object->entity && object->entity->className == "light")
                count++;
        }

        size_t& count;
    };

    struct StartPointFinder
    {
        StartPointFinder(ReversibleTransform &transform)
            : transform(transform) {}

        void operator()(boost::shared_ptr<Object> object)
        {
            if(object->entity)
            {
                const std::string &className = object->entity->className;
                if(className == "info_player_start" || className == "info.player.start")
                {
                    OrthoTransform trans;
                    trans.SetVector(-object->transform.GetVector());
                    trans.SetMatrix(object->transform.GetMatrix());
                    transform = trans;
                }
            }
        }

        ReversibleTransform &transform;
    };

    struct BrushEncoder
    {
        BrushEncoder(short *brushBuffer, GpuFace *faceBuffer, size_t &brushPosition, size_t &facePosition)
            : brushBuffer(brushBuffer), faceBuffer(faceBuffer), brushPosition(brushPosition), facePosition(facePosition) {}

        void operator()(boost::shared_ptr<Object> object)
        {
            boost::shared_ptr<Brush> brush = object->brush;
            if(!brush)
                return;

            // Write the brush offset and the number of faces.
            brushBuffer[brushPosition++] = facePosition;
            brushBuffer[brushPosition++] = facePosition + brush->faces.size();
            brushPosition += 2; // Skip unused

            // Write each one of the faces.
            for(size_t i = 0; i < brush->faces.size(); i++)
            {
                const BrushFace &face = brush->faces[i];

                faceBuffer[facePosition].plane = face.plane;

                // Advance the face position.
                facePosition++;
            }
        }

       short *brushBuffer;
       GpuFace *faceBuffer;
       size_t &brushPosition;
       size_t &facePosition;
    };

    struct LightEncoder
    {
        LightEncoder(GpuLight *lightBuffer, size_t &lightPosition)
            : lightBuffer(lightBuffer), lightPosition(lightPosition) {}

        void operator()(boost::shared_ptr<Object> object)
        {
            if(!object->entity || object->entity->className != "light")
                return;

            lightBuffer[lightPosition++].position = object->transform.GetVector();
        }

        GpuLight *lightBuffer;
        size_t &lightPosition;
    };

    void RayTracer::InitScene(boost::shared_ptr<Object> object)
    {
        // Initialize the opengl matrices.
        glMatrixMode(GL_MODELVIEW);
        glLoadIdentity();

        glMatrixMode(GL_PROJECTION);
        glLoadIdentity();
        glOrtho(0.0, 1.0, 0.0, 1.0, -1.0, 1.0);

        // Count the number of faces.
        faceCount = 0;
        object->Traverse(FaceCounter(faceCount));

        // Count the number of brushes.
        brushCount = 0;
        object->Traverse(BrushCounter(brushCount));

        // Count the number of lights.
        lightCount = 0;
        object->Traverse(LightCounter(lightCount));

        // Find the start point.
        object->Traverse(StartPointFinder(cameraTransform));

        // Build the gpu data structures.

        // Create the brush buffers.
        faceBuffer.reset(new LinearBuffer(faceCount*sizeof(GpuFace), GL_RGBA32F_ARB));
        brushBuffer.reset(new LinearBuffer(brushCount*8, GL_RGBA16I_EXT));

        // Populate the brush buffers.
        char *faceTempBuffer = new char[faceBuffer->GetSize()];
        char *brushTempBuffer = new char[brushBuffer->GetSize()];

        size_t brushPosition = 0;
        size_t facePosition = 0;

        object->Traverse(BrushEncoder((short*)brushTempBuffer, (GpuFace*)faceTempBuffer, brushPosition, facePosition));

        faceBuffer->Write(faceTempBuffer, 0, faceBuffer->GetSize());
        brushBuffer->Write(brushTempBuffer, 0, brushBuffer->GetSize());
        
        delete [] faceTempBuffer;
        delete [] brushTempBuffer;

        // Create the light buffer
        lightBuffer.reset(new LinearBuffer(lightCount*sizeof(GpuLight), GL_RGBA32F_ARB));

        // Populate the light buffer.
        char *tempLightBuffer = new char[lightBuffer->GetSize()];

        size_t lightPosition = 0;

        object->Traverse(LightEncoder((GpuLight*)tempLightBuffer, lightPosition));

        lightBuffer->Write(tempLightBuffer, 0, lightBuffer->GetSize());

        delete [] tempLightBuffer;

        // Create a jitter buffer.
        jitterSize = 256;
        jitterBuffer.reset(new LinearBuffer(jitterSize*sizeof(float)*4, GL_RGBA32F_ARB));

        // Populate the jitter buffer.
        float *tempJitterBuffer = new float[jitterSize*4];

        srand(time(NULL));
        for(size_t i = 0; i < jitterSize*4; i++)
        {
            tempJitterBuffer[i] = float(rand())/float(RAND_MAX);
        }
        
        jitterBuffer->Write(tempJitterBuffer, 0, jitterBuffer->GetSize());
        delete [] tempJitterBuffer;

        // Create the render state buffer.
        framebuffer.reset(new TargetBuffer(width, height));
        colorBuffer.reset(new PingPongBuffer(framebuffer, 0,
            boost::make_shared<TextureBuffer> (width, height, 128, GL_RGBA32F_ARB, GL_RGBA, GL_FLOAT),
            boost::make_shared<TextureBuffer> (width, height, 128, GL_RGBA32F_ARB, GL_RGBA, GL_FLOAT)));
        rayBuffer.reset(new PingPongBuffer(framebuffer, 1,
            boost::make_shared<TextureBuffer> (width, height, 128, GL_RGBA32F_ARB, GL_RGBA, GL_FLOAT),
            boost::make_shared<TextureBuffer> (width, height, 128, GL_RGBA32F_ARB, GL_RGBA, GL_FLOAT)));
        normalBuffer.reset(new PingPongBuffer(framebuffer, 2,
            boost::make_shared<TextureBuffer> (width, height, 128, GL_RGBA32F_ARB, GL_RGBA, GL_FLOAT),
            boost::make_shared<TextureBuffer> (width, height, 128, GL_RGBA32F_ARB, GL_RGBA, GL_FLOAT)));

        // Create the primary rays buffer.
        primaryRays.reset(new TextureBuffer(width, height, 128, GL_RGBA32F_ARB, GL_RGBA, GL_FLOAT));

        // Create the ray screen.
        float *rayDirection = new float[width*height*4];
        float *rayOrigin = new float[width*height*4];

        float screenWidth = 40.0f;
        float screenHeight = 30.0f;

        Vector3 dx = Vector3(screenWidth/float(width), 0.0f, 0.0f);
        Vector3 dy = Vector3(0.0f, screenHeight/float(height), 0.0f);
        Vector3 bottomLeft(-screenWidth*0.5f, -screenHeight*0.5f, 0.0f);
        Vector3 origin(0.0f, 0.0f, 10.0f);

        dx = cameraTransform.GetInverseMatrix() * dx;
        dy = cameraTransform.GetInverseMatrix() * dy;
        bottomLeft = bottomLeft / cameraTransform;
        origin = origin / cameraTransform;

        /*printf("dx %f %f %f\n", dx.x, dx.y, dx.z);
        printf("dy %f %f %f\n", dy.x, dy.y, dy.z);
        printf("bottomLeft %f %f %f\n", bottomLeft.x, bottomLeft.y, bottomLeft.z);
        printf("origin %f %f %f\n", origin.x, origin.y, origin.z);*/

        for(int y = 0; y < height; y++)
        {
            for(int x = 0; x < width; x++)
            {
                Vector3 screenPos = bottomLeft + x*dx + y*dy;
                Vector3 direction = (screenPos - origin).Normalize();

                int offset = y*width + x;

                rayOrigin[offset*4] = origin.x;
                rayOrigin[offset*4+1] = origin.y;
                rayOrigin[offset*4+2] = origin.z;
                rayOrigin[offset*4+3] = 0.0f;

                rayDirection[offset*4] = direction.x;
                rayDirection[offset*4+1] = direction.y;
                rayDirection[offset*4+2] = direction.z;

                //printf("Direction: %f %f %f\n", direction.x, direction.y, direction.z);
            }
        }

        rayBuffer->Write(rayOrigin, 0, rayBuffer->GetSize());
        primaryRays->Write(rayDirection, 0, primaryRays->GetSize());

        delete [] rayDirection;
        delete [] rayOrigin;
    }

    void RayTracer::DisplayScene()
    {
        //valid = false;
        //cameraTransform.SetVector(cameraTransform.GetVector() + Vector3(0.0, 0.0, 5.f));

        if(!valid)
            RenderScene();

        if(!valid)
            return;

        // Draw to the back buffer.
        glBindFramebufferEXT(GL_FRAMEBUFFER_EXT, 0);
        glDrawBuffer(GL_BACK);
        glClear(GL_COLOR_BUFFER_BIT);

        // Activate the passthrough program.
        passthroughProgram->Activate();

        // Setup the program.
        colorBuffer->Activate(0);
        passthroughProgram->SetUniform("image", 0);

        // Draw the screen quad.
        ScreenQuad();
    }

    void RayTracer::RenderScene()
    {
        // Activate the target buffer.
        framebuffer->Activate();
        GLenum drawBuffers[] = {
            GL_COLOR_ATTACHMENT0_EXT,
            GL_COLOR_ATTACHMENT1_EXT,
            GL_COLOR_ATTACHMENT2_EXT,
        };

        glDrawBuffers(3, drawBuffers);

        // Cast the primary rays.
        CastPrimaryRays();

        // Accumulate shading.
        AccumShade();

        valid = true;
    }

    void RayTracer::CastPrimaryRays()
    {
        // Activate the primary ray casting program.
        primaryCastProgram->Activate();

        // Setup primary ray casting.
        rayBuffer->Activate(0);
        primaryRays->Activate(1);
        faceBuffer->Activate(2);
        brushBuffer->Activate(3);

        primaryCastProgram->SetUniform("rayOrigins", 0);
        primaryCastProgram->SetUniform("rayDirections", 1);
        primaryCastProgram->SetUniform("brushFaces", 2);
        primaryCastProgram->SetUniform("brushes", 3);

        primaryCastProgram->SetUniform("brushCount", brushCount);
        primaryCastProgram->SetUniform("faceCount", faceCount);

        // Ray cast the primary rays.
        ScreenQuad();

        // Change the render buffer output.
        colorBuffer->Pong();
        rayBuffer->Pong();
        normalBuffer->Pong();
    }

    void RayTracer::AccumShade()
    {
        // Activate the shade program.
        shadeProgram->Activate();

        // Setup shading.
        colorBuffer->Activate(0);
        rayBuffer->Activate(1);
        normalBuffer->Activate(2);

        faceBuffer->Activate(3);
        brushBuffer->Activate(4);
        lightBuffer->Activate(5);
        jitterBuffer->Activate(6);

        shadeProgram->SetUniform("prevColors", 0);
        shadeProgram->SetUniform("prevPositions", 1);
        shadeProgram->SetUniform("prevNormals", 2);

        shadeProgram->SetUniform("brushFaces", 3);
        shadeProgram->SetUniform("brushes", 4);
        shadeProgram->SetUniform("lights", 5);
        shadeProgram->SetUniform("jitterBuffer", 6);

        shadeProgram->SetUniform("brushCount", brushCount);
        shadeProgram->SetUniform("faceCount", faceCount);
        shadeProgram->SetUniform("lightCount", lightCount);
        shadeProgram->SetUniform("jitterSize", jitterSize);

        // Shade the current position.
        ScreenQuad();

        // Change the render buffer output.
        colorBuffer->Pong();
        rayBuffer->Pong();
        normalBuffer->Pong();
    }

    void RayTracer::ScreenQuad()
    {
        glBegin(GL_TRIANGLE_STRIP); // GL_QUADS is deprecated.
        glTexCoord2f(0.0f, 0.0f); glVertex3f(0.0f, 0.0f, 0.0f);
        glTexCoord2f(1.0f, 0.0f); glVertex3f(1.0f, 0.0f, 0.0f);
        glTexCoord2f(0.0f, 1.0f); glVertex3f(0.0f, 1.0f, 0.0f);
        glTexCoord2f(1.0f, 1.0f); glVertex3f(1.0f, 1.0f, 0.0f);
        glEnd();
    }
};

